@karen-sy this looks great!

I just played around with a bit and tried to add a visual to view the pose grids by displaying "axis arrows" at each grid point. Currently, it looks like the original pose is not at the center of the pose grid, but at the bottom of it. Could you please take a look at this and check if this is right, or if my visualization is erroneous? (It may well be...) If this is right, can you please change the gridding code to ensure the original pose is at the grid center?

Also, it also looks like when the rotation grid has only 1 grid point, it ends up being rotated relative to the original pose, rather than having the same rotation to the original pose. This may again be an issue with my visualization, but if it is a real property of the pose grid, I think it would be better to ensure that with small rotation grids, we cover the original rotation exactly.

Also, it also looks like when the rotation grid has only 1 grid point, it ends up being rotated relative to the original pose, rather than having the same rotation to the original pose. This may again be an issue with my visualization, but if it is a real property of the pose grid, I think it would be better to ensure that with small rotation grids, we cover the original rotation exactly.

This is likely from default jnp.linspace behavior -- I'll modify the grid ordering so that we order the close points (to center) first. Re: viz, let me check again

I modified the gridding API such that every grid is guaranteed to have the original pose included:

def pose_grid(
    pose_center: Pose,   # center pose (will always be in returned Pose)
    half_dx=float,          # HALF dimension size of grid (i.e. if we want to grid over range [-n, n] for n>=0, then half_dx = n)
    half_dy=float,
    half_dz=float,
    half_nx=int,             # number of gridding points in HALF dimension (i.e. if we want 2k points in a grid of [-n,n], half_nx = k)
    half_ny=int,
    half_nz=int,
    half_dangle=float,
    half_n_xrot=int,       
    half_n_yrot=int,
    half_n_zrot=int,
) -> Pose

Perf benchmark: Time taken: 0.6334781646728516 milliseconds for 15625 poses

Also, not exactly sure what you were exactly visualizing; could you toggle VIZ_TEST on L168 to True and see the mug + axes that are visualized on rerun? (@georgematheos )

Oops somehow I thought you merged the branch already! Thanks so much for doing the test refactoring. Merging!